
from pydantic import Field
from core.llm.base_analyze import BaseAnalyzer
from core import prompt_dir
from typing import Dict, List
import pandas as pd
from loguru import logger

class DatasetUsageAnalyzer(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / '数据表描述.md')

    def invoke(
            self,
            schema: Dict = Field(description="数据表的结构信息"),
            rows: List[Dict] = Field(description="数据行")
    ):
        df = pd.DataFrame(rows)
        prompt = self.template_prompt.format(
            inputs={
                "schema": schema,
                "preview": df.head(50).to_markdown(),
            },
            remove_template_variables=True
        )
        data = self.analyze(prompt)
        if not isinstance(data, dict):
            return {}
        elif 'description' not in data:
            return {}
        elif 'fields' not in data:
            return {}
        if 'query' in schema:
            data['query'] = schema['query']
        data['name'] = schema.get('name', '')
        data['fields'] = str(data['fields'])
        data['columns'] = schema['columns']
        data['rows'] = rows
        return data
